package com.eobard.obfuscate.optional;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.file.FileNameUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.eobard.config.obfuscate.ObfuscatorConfig;
import com.eobard.enums.PyDependenciesEnum;
import com.eobard.enums.PyInstallEnum;
import com.eobard.obfuscate.IObfuscator;
import com.eobard.params.MnnConverterParams;
import com.eobard.utils.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.awt.*;
import java.io.File;
import java.nio.file.Paths;
import java.util.List;
import java.util.function.BiFunction;
import java.util.function.Function;

/**
 * @Author: Eobard Gu
 * @Date: 2025-09-30
 * @Description: Mnn模型转换器核心类
 */
@Component
@Slf4j
public class MnnConverterObfuscator implements IObfuscator {

    @Value("${obfuscate.enable.typeSolver.logging:false}")
    private boolean print;


    //mnn模型文件的路径位置
    private static String mnnModelPath;

    //MNN转换库的安装位置根路径
    private static String mnnInstallRoot;

    //mnn->json 和 json->mnn的中间临时保存位置
    private static String converterPath;

    //指定修改md5的模型文件列表，默认修改全部
    private static List<String> mnnModelNames = null;

    //指定为mnn->json的过程
    private static boolean forwardConverter = true;

    //pip安装是否使用默认命令行
    private static boolean pipUsingDefault = false;

    private void sourceCodeInit(MnnConverterParams params) {
        //MNN库安装位置检查
        String installRoot = params.getMnnInstallRoot();
        ExceptionUtils.checkNonNull(installRoot, "mnn源码安装位置不能为空");
        mnnInstallRoot = installRoot.endsWith("/build") ? installRoot : installRoot + "/build";
        EnvUtils.checkEnv(mnnInstallRoot + MnnUtils.MNN_VERSION);
    }

    private void pipInit() {
        // 检查解释器是否存在
        EnvUtils.isPythonAvailable(PythonUtils.PYTHON3_EXPLAINER);
        //检查依赖是否存在
        EnvUtils.hasPythonDependencies(PyDependenciesEnum.MNN.dependencies);

        //判断是否可以使用mnn命令
        try {
            EnvUtils.run(MnnUtils.MNN_VERSION_PYTHON);
        } catch (Exception e) {
            log.info("\n command not found:【 {} 】, trying to use 【 {} 】", MnnUtils.MNN_VERSION_PYTHON, MnnUtils.MNN_VERSION_PYTHON_DEFAULT);
            try {
                EnvUtils.run(MnnUtils.MNN_VERSION_PYTHON_DEFAULT);
                pipUsingDefault = true;
            } catch (Exception e1) {
                throw new RuntimeException(e1);
            }
        }
    }


    //初始化数据和检查
    @Override
    public void initialize(ObfuscatorConfig config) {
        MnnConverterParams params = config.getMnnConverterParams();

        //如果是源码安装MNN库
        if (params.getInstallEnum() == PyInstallEnum.SOURCE_CODE) {
            sourceCodeInit(params);
        } else {
            //否则是pip安装
            pipInit();
        }

        //获取正向还是反向转换(正:mnn->json , 反:json->mnn)
        forwardConverter = params.isForwardConverter();

        //项目模型路径检查
        mnnModelPath = params.getMnnModelPath();
        //判断mnnResourceRoot是否存在.mnn结尾的模型文件
        FileUtils.checkFileExt(mnnModelPath, FileUtils.MODEL);
        //获取所有以.mnn结尾的模型文件
        List<String> mnnFiles = FileUtils.loopTypeExtFiles(mnnModelPath, MnnUtils.MNN_EXT);


        //指定转换mnn模型名称检查:如果modelNames为空,默认修改mnnModelPath所有模型文件,否则就修改指定名称模型文件
        List<String> modelNames = params.getMnnModelNames();
        if (CollectionUtil.isEmpty(modelNames) || modelNames.isEmpty()) {
            mnnModelNames = mnnFiles;
        } else {
            mnnModelNames = modelNames.stream()
                    .map(name -> name.startsWith(mnnModelPath) ? name : Paths.get(mnnModelPath, name).toString())
                    .filter(mnnFiles::contains).toList();
        }

        //创建临时转换目录
        if (StrUtil.isBlank(params.getConverterPath())) {
            converterPath = System.getProperty("user.home") + "/mnn_converter_output";
        } else {
            converterPath = params.getConverterPath();
        }
        File file = new File(converterPath);
        if (!file.exists()) {
            file.mkdirs();
        }
    }

    //混淆核心代码
    @Override
    public void obfuscate(ObfuscatorConfig config) {
        if (forwardConverter) {
            convertModels(
                    mnnModelNames,
                    mnn -> converterPath + File.separator + FileNameUtil.getName(mnn).replace(MnnUtils.MNN_EXT, MnnUtils.JSON_EXT),
                    (input, output) -> config.getMnnConverterParams().getInstallEnum() == PyInstallEnum.SOURCE_CODE ?
                            MnnUtils.mnnToJsonCommand(input, output) : MnnUtils.mnnToJsonPythonCommand(input, output, pipUsingDefault),
                    config.getMnnConverterParams().getInstallEnum()
            );
        } else {
            List<String> jsonFiles = FileUtils.loopTypeExtFiles(converterPath, MnnUtils.JSON_EXT);
            convertModels(
                    jsonFiles,
                    json -> {
                        String newName = FileNameUtil.mainName(json) + "_" + String.valueOf(System.currentTimeMillis()).substring(7) + MnnUtils.MNN_EXT;
                        return json.replace(FileNameUtil.getName(json), newName);
                    },
                    (input, output) -> config.getMnnConverterParams().getInstallEnum() == PyInstallEnum.SOURCE_CODE ?
                            MnnUtils.jsonToMnnCommand(input, output) : MnnUtils.jsonToMnnPythonCommand(input, output, pipUsingDefault),
                    config.getMnnConverterParams().getInstallEnum()
            );
        }
    }

    //模型转换
    private void convertModels(List<String> inputFiles,
                               Function<String, String> outputNameGenerator,
                               BiFunction<String, String, String> commandGenerator,
                               PyInstallEnum installEnum) {

        inputFiles.forEach(input -> {
            String output = outputNameGenerator.apply(input);
            //创建命令
            String command = installEnum == PyInstallEnum.SOURCE_CODE ?
                    mnnInstallRoot + commandGenerator.apply(input, output) : commandGenerator.apply(input, output);

            if (print) {
                log.info("\n current MNN env is 【 {} 】,using【 {} 】", installEnum, command);
            }
            //执行命令
            EnvUtils.run(command);
        });

        //打开输出的文件夹
        try {
            System.setProperty("java.awt.headless", "false");
            Desktop.getDesktop().open(new File(converterPath));
        } catch (Exception e) {
            log.error("当前环境不支持桌面操作: " + converterPath);
        }
    }

    //是否启用这个功能
    @Override
    public boolean supports(ObfuscatorConfig config) {
        return ObjectUtil.isNotEmpty(config.getMnnConverterParams());
    }
}
